import os
import time

import cv2
import numpy as np
import shutil
import random

try:
    from .cvio import cvio
    from . import imagesize
except:
    from cvio import cvio
    import imagesize

def remove_all_files_by_ext(src, ext_type='xml'):
    files = cvio.load_ext_list(src, ext_type=ext_type, recursive=True)
    print(len(files))
    for file in files:
        os.remove(file)

def remove_augmented_images(src):
    imgann = cvio.load_img_ann_list(src, recursive=True)
    print(len(imgann))
    todel = []
    c = 0
    for i, a in imgann:
        if ('_rotate' in i.lower() or '_L' in i
            or '_R' in i or 'mosaic' in i.lower()
            or '_flip' in i.lower() or 'left' in i.lower()
                or 'right' in i.lower() or 'flip' in i.lower()):
            todel.append([i, a])
            c += 1
            os.remove(i)
            os.remove(a)
    print(c, len(imgann))


def remove_images_without_annotations(src):
    imgs = cvio.load_image_list(src, recursive=True)
    tbar = enumerate(imgs, 1)
    c = 0
    n = len(imgs)
    for i, img in tbar:
        ann = os.path.splitext(img)[0] + '.json'
        if not os.path.exists(ann):
            c += 1
            os.remove(img)
        print('[%d/%d] %d' % (i, n, c))


def refine_and_filter_all_json(src):
    anns = cvio.load_ext_list(src, ext_type='json', recursive=True)
    c = 0
    tbar = enumerate(anns, 1)
    n = len(anns)
    for i, ann in tbar:
        try:
            annInfo = cvio.load_ann(ann)
            if not 'imagePath' in annInfo:
                c += 1
                print(ann)
                os.remove(ann)
                continue
            annInfo['imagePath'] = os.path.basename(annInfo['imagePath'])
            annInfo['imageData'] = None
            cvio.write_ann(annInfo, ann)
        except:
            c += 1
            os.remove(ann)
            print(ann)
    print('[%d/%d] exceptions %d' % (i, n, c))


def write_all_images_name(src):
    imgs = cvio.load_image_list(src, recursive=True)
    tbar = enumerate(imgs, 1)
    c = 0
    n = len(imgs)
    fp = open(os.path.join(src, 'imageList.txt'), 'w')
    for i, img in tbar:
        print('[%d/%d]' % (i,n), img, file=fp)
    fp.close()


def remove_all_unneed_labels(src):
    imgann = cvio.load_img_ann_list(src, recursive=True)
    n = len(imgann)
    print()
    c = 0
    d = 0
    for img, ann in imgann:
        annInfo = cvio.load_ann(ann)
        shapes = []
        for s in annInfo['shapes']:
            k = s['label']
            if 'daiding' in k or k.startswith('sp_'):  # or 'fake' in k
                continue
            shapes.append(s)
        if not len(shapes):
            d += 1
            os.remove(img)
            os.remove(ann)
            continue
        if len(shapes) != len(annInfo['shapes']):
            c += 1
            annInfo['shapes'] = shapes
            annInfo['imageData'] = None
            cvio.write_ann(annInfo, ann)
    print(d, c, n)


def linear_transform(img, a, b):
    '''
    :param img: [h, w, 3] 彩色图像
    :param a:  float  这里需要是浮点数，把图片uint8类型的数据强制转成float64
    :param b:  float
    :return: out = a * img + b
    '''
    out = a * img + b
    out[out > 255] = 255
    out = np.around(out)
    out = out.astype(np.uint8)
    return out


def rename_dataset(src, shuffle=True, prefix=''):
    imgs = cvio.load_image_list(src,recursive=False,silent=False)
    if shuffle:
        random.shuffle(imgs)
    n = len(imgs)
    for i,p in enumerate(imgs,1):
        b, e = os.path.splitext(p)
        # _time = time.strftime("%Y%m%d%H%M%S", time.localtime())
        # d = _time[:8]
        # t = _time[8:]
        name = "%s_%d%s" % (prefix, i,e)
        print("[%d/%d]" % (i,n), os.path.basename(p), name)
        dst = os.path.join(os.path.dirname(p), name)
        os.rename(p, dst)
        ann = b + '.json'
        if not os.path.exists(ann):
            continue
        try:
            ast = os.path.splitext(dst)[0] + '.json'
            os.rename(ann, ast)
            annInfo = cvio.load_ann(ast)
            annInfo['imagePath'] = name
            annInfo['imageData'] = None
            wh = imagesize.get(dst)
            annInfo["imageWidth"] = wh[0]
            annInfo["imageHeight"] = wh[1]
            cvio.write_ann(annInfo, ast)
        except:
            print("标注文件异常", ann)


def check_unmatch_imgs_between_cv_pil(src):
    imgs = cvio.load_image_list(src, recursive=True)
    c = 0
    tbar = enumerate(imgs, 1)
    n = len(imgs)
    for i, img in tbar:
        print('%d/%d %d %s' % (i, n, c, img))
        cvimg = cvio.load_img(img)
        pilimg = cvio.load_img(img, mode='pil')
        cvh, cvw = cvimg.shape[:-1]
        pilw, pilh = pilimg.size
        if cvh == pilh and cvw == pilw:
            continue
        elif cvh != pilw and cvw != pilh:
            print('something unexpected errors!')
            continue
        else:
            c += 1
            cvio.write_img(np.array(cvimg, dtype=np.uint8), img)

def collect_images_with_multi_labels(src):
    imglist = cvio.load_image_list(src, recursive=True, silent=False)
    cls2imgs = {}
    for img in imglist:
        label = img.split(os.sep)[-2]
        if not label in cls2imgs:
            cls2imgs[label] = [img]
        else:
            cls2imgs[label] += [img]
            
    
    for label, imgs in cls2imgs.items():
        imgdic = {}
        for img in imgs:
            if '_points' in img:
                bn = os.path.basename(img).split('_points')[0]
            else:
                n, e = os.path.splitext(os.path.basename(img))
                ns = n.split('_')[:-1]
                n = ''
                for x in ns:
                    n += '%s_' % x
                n = n[:-1]
                bn = '%s%s' % (n, e)

            if not bn in imgdic:
                imgdic[bn] = [img]
            else:
                imgdic[bn].append(img)
        
        images = []
        for bn, imgs in imgdic.items(): 
            if len(imgs) <= 1: continue
            images.extend(imgs)
        if not len(images): continue
        print(label, len(images))
        dst = os.path.join(src, label, 'repeat')
        if not os.path.exists(dst):
            os.makedirs(dst)
        for img in images:
            shutil.move(img, os.path.join(dst, os.path.basename(img)))
    
def check_cuts_and_zb_labels(cutspath, zbtxt):
    lc1 = [f for f in os.listdir(cutspath) if os.path.isdir(os.path.join(cutspath, f)) and not '总表中不存在的标签' in f]
    with open(zbtxt) as fp:
        lc2 = [f.strip() for f in fp.readlines()]
    savepath = os.path.join(cutspath, '总表中不存在的标签')
    if not os.path.exists(savepath):
        os.makedirs(savepath)
    tbar = set(lc1).difference(set(lc2))
    if not len(tbar):
        print('标签完全匹配（%s）' % cutspath)
        return
    print('总表中不存在以下标签:')
    i = 0
    for label in tbar:
        i += 1
        save = os.path.join(savepath, label)
        print(i, label, '移动至\n%s.' % save)
        shutil.move(os.path.join(cutspath, label), save)

def resize_dataset(src, dst='', with_bbox=True, limit=1333):
    images = cvio.load_image_list(src, silent=False)
    n = len(images)
    if dst == '' or dst is None:
        dst = os.path.abspath(os.path.join(src, '..', 'compress'))
    if dst != '' and not os.path.exists(dst):
        os.makedirs(dst)
    for i, imgpath in enumerate(images, 1):
        try:
            img = np.array(cvio.load_img(imgpath, mode='pil'))
            h, w = img.shape[:2]
            max_edge = max(h, w)
            if max_edge > limit:
                scale_factor = limit / max_edge
                height = int(round(h * scale_factor))
                width = int(round(w * scale_factor))
                img = cv2.resize(img, (width, height), interpolation=cv2.INTER_LINEAR)
                dstimg = imgpath if dst == '' else os.path.join(dst, os.path.basename(imgpath))
                cvio.write_img(img, dstimg, mode='pil')
                if with_bbox:
                    ann = os.path.splitext(imgpath)[0] + '.json'
                    if os.path.exists(ann):
                        ann_info = cvio.load_ann(ann)
                        ann_info['imageWidth'] = width
                        ann_info['imageHeight'] = height
                        for j, shape in enumerate(ann_info['shapes']):
                            points = shape['points']
                            shape['points'] = [[x[0] * scale_factor, x[1] * scale_factor] for x in points]
                        dstann = ann if dst == '' else os.path.join(dst, os.path.basename(ann))
                        cvio.write_ann(ann_info, dstann)
            print('[%d/%d] %s %s %s' % (i, n, str((h, w)), str(img.shape[:2]), os.path.basename(imgpath)))
        except:
            print('Exception (%s).' % imgpath)
